import pandas as pd


# 分组.py是对个电影类型分类的程序,把各个小类合并成一个大类
def f(x):
    return 1 if x > 0 else 0


df = pd.read_csv(r"DSproject/csv/preprocessed.csv")

# 使用 apply 逐行计算
df["joy_and_escapism"] = (
    df[
        [
            "is_Comedy",
            "is_Romance",
            "is_Animation",
            "is_Musical",
            "is_Music",
            "is_Fantasy",
            "is_Family",
        ]
    ]
    .sum(axis=1)
    .apply(f)
)
df["Stimulation_and_Adventure"] = (
    df[["is_Action", "is_Adventure", "is_Sci-Fi", "is_Sport"]].sum(axis=1).apply(f)
)
df["Realism_and_Depth"] = (
    df[["is_Drama", "is_Biography", "is_History"]].sum(axis=1).apply(f)
)
df["Anxiety_Conflict"] = (
    df[["is_War", "is_Crime", "is_Thriller", "is_Mystery", "is_Film-Noir", "is_Horror"]]
    .sum(axis=1)
    .apply(f)
)
df["Documentary_and_Public"] = (
    df[["is_Documentary", "is_Reality-TV", "is_Talk-Show"]].sum(axis=1).apply(f)
)
df["Other"] = df[["is_Short", "is_Western"]].sum(axis=1).apply(f)

df.to_csv(r"DSProject\csv\preprocessed2.csv", index=False)
